#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
该脚本用于人体姿态估计，基于OpenPose模型。支持三种不同的模型：BODY25、COCO 和 MPI。
BODY_25模型（--model_pose BODY_25）包括身体和脚部关键点，
COCO 是一个大规模的目标检测、分割和人体关键点检测数据集。17 个关键点
MPII 是一个广泛使用的人体姿态估计数据集，主要关注日常活动。16个关键点
OCHuman 是一个包含复杂场景和遮挡情况的人体姿态估计数据集。
CrowdPose 是一个专门针对拥挤场景的人体姿态估计数据集。
PoseTrack 是一个多目标跟踪和人体姿态估计数据集，主要用于视频中的多目标跟踪。
"""

from __future__ import division
import cv2
import time
import numpy as np
import matplotlib.pyplot as plt
import os

class general_pose_model(object):
    """
    通用姿态估计模型类，支持多种姿态估计模型（BODY25, COCO, MPI）。
    """

    def __init__(self, modelpath, mode="BODY25"):
        """
        初始化姿态估计模型。

        参数:
            modelpath (str): 模型文件路径。
            mode (str): 使用的姿态估计模型类型，默认为 "BODY25"。
        """
        self.inWidth = 368   # 输入图像宽度
        self.inHeight = 368  # 输入图像高度
        self.threshold = 0.1  # 置信度阈值

        if mode == "BODY25":
            self.pose_net = self.general_body25_model(modelpath)
        elif mode == "COCO":
            self.pose_net = self.general_coco_model(modelpath)
        elif mode == "MPI":
            self.pose_net = self.get_mpi_model(modelpath)
    
    def get_hald_model(self, modelpath):
        """
        加载 MPI 模型。

        参数:
            modelpath (str): 模型文件路径。

        返回:
            mpi_model: 加载的 MPI 模型。
        """
        self.points_name = {
            "Head": 0, "Neck": 1,
            "RShoulder": 2, "RElbow": 3, "RWrist": 4,
            "LShoulder": 5, "LElbow": 6, "LWrist": 7,
            "RHip": 8, "RKnee": 9, "RAnkle": 10,
            "LHip": 11, "LKnee": 12, "LAnkle": 13,
            "Chest": 14, "Background": 15
        }
        self.num_points = 22
        self.point_pairs =  [ [0,1],[1,2],[2,3],[3,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[10,11],[11,12],[0,13],[13,14],[14,15],[15,16],[0,17],[17,18],[18,19],[19,20] ]  # 关键点连接对

        prototxt = os.path.join(
            modelpath,
            "hand/pose_deploy.prototxt"
        )
        caffemodel = os.path.join(
            modelpath,
            "hand/pose_iter_102000.caffemodel"
        )
        mpi_model = cv2.dnn.readNetFromCaffe(prototxt, caffemodel)
        return mpi_model
    def get_mpi_model(self, modelpath):
        """
        加载 MPI 模型。

        参数:
            modelpath (str): 模型文件路径。

        返回:
            mpi_model: 加载的 MPI 模型。
        """
        self.points_name = {
            "Head": 0, "Neck": 1,
            "RShoulder": 2, "RElbow": 3, "RWrist": 4,
            "LShoulder": 5, "LElbow": 6, "LWrist": 7,
            "RHip": 8, "RKnee": 9, "RAnkle": 10,
            "LHip": 11, "LKnee": 12, "LAnkle": 13,
            "Chest": 14, "Background": 15
        }
        self.num_points = 15
        self.point_pairs = [
            [0, 1], [1, 2], [2, 3],
            [3, 4], [1, 5], [5, 6],
            [6, 7], [1, 14], [14, 8],
            [8, 9], [9, 10], [14, 11],
            [11, 12], [12, 13]
        ]
        prototxt = os.path.join(
            modelpath,
            "pose/mpi/pose_deploy_linevec_faster_4_stages.prototxt"
        )
        caffemodel = os.path.join(
            modelpath,
            "pose/mpi/pose_iter_160000.caffemodel"
        )
        mpi_model = cv2.dnn.readNetFromCaffe(prototxt, caffemodel)
        return mpi_model

    def general_coco_model(self, modelpath):
        """
        加载 COCO 模型。

        参数:
            modelpath (str): 模型文件路径。

        返回:
            coco_model: 加载的 COCO 模型。
        """
        self.points_name = {
            "Nose": 0, "Neck": 1,
            "RShoulder": 2, "RElbow": 3, "RWrist": 4,
            "LShoulder": 5, "LElbow": 6, "LWrist": 7,
            "RHip": 8, "RKnee": 9, "RAnkle": 10,
            "LHip": 11, "LKnee": 12, "LAnkle": 13,
            "REye": 14, "LEye": 15,
            "REar": 16, "LEar": 17,
            "Background": 18
        }
        self.num_points = 18
        self.point_pairs = [
            [1, 0], [1, 2], [1, 5],
            [2, 3], [3, 4], [5, 6],
            [6, 7], [1, 8], [8, 9],
            [9, 10], [1, 11], [11, 12],
            [12, 13], [0, 14], [0, 15],
            [14, 16], [15, 17]
        ]
        prototxt = os.path.join(
            modelpath,
            "pose/coco/pose_deploy_linevec.prototxt"
        )
        caffemodel = os.path.join(
            modelpath,
            "pose/coco/pose_iter_440000.caffemodel"
        )
        coco_model = cv2.dnn.readNetFromCaffe(prototxt, caffemodel)
        return coco_model

    def general_body25_model(self, modelpath):
        """
        加载 BODY25 模型。

        参数:
            modelpath (str): 模型文件路径。

        返回:
            body25_model: 加载的 BODY25 模型。
        """

        """
        模型的关键点连接对（point_pairs），用于表示人体各个部位之间的连接关系
        self.num_points = 25：指定了该模型共有 25 个关键点（即人体的 25 个特定部位）。
        self.point_pairs：是一个二维列表，每个子列表包含两个整数，表示两个关键点之间的连接关系
        。例如 [1, 0] 表示第 1 个关键点和第 0 个关键点之间有连接。
        """
        self.num_points = 25
        self.point_pairs = [
            [1, 0], [1, 2], [1, 5],
            [2, 3], [3, 4], [5, 6],
            [6, 7], [0, 15], [15, 17],
            [0, 16], [16, 18], [1, 8],
            [8, 9], [9, 10], [10, 11],
            [11, 22], [22, 23], [11, 24],
            [8, 12], [12, 13], [13, 14],
            [14, 19], [19, 20], [14, 21]
        ]
        prototxt = os.path.join(
            modelpath,
            "pose/body_25/pose_deploy.prototxt"
        )
        caffemodel = os.path.join(
            modelpath,
            "pose/body_25/pose_iter_584000.caffemodel"
        )
        # print(cv2)
        body25_model = cv2.dnn.readNetFromCaffe(prototxt, caffemodel)
        return body25_model
    def predict(self, img_cv2):
        """
        对输入图像进行姿态估计预测。

        参数:
            imgfile (str): 输入图像文件路径。

        返回:
            points (list): 关键点坐标列表。
        """
        # img_cv2 = cv2.imread(imgfile)
        img_height, img_width, _ = img_cv2.shape
        """
        使用 cv2.dnn.blobFromImage 函数将输入图像转换为适合神经网络处理的格式（Blob）。
1.0 / 255：将图像像素值从 [0, 255] 归一化到 [0, 1]。
(self.inWidth, self.inHeight)：指定输入图像的宽度和高度（默认为 368x368）。
(0, 0, 0)：均值减去操作，这里没有实际减去任何值。
swapRB=False：不交换 BGR 和 RGB 通道顺序。
crop=False：不裁剪图像。
"""
        inpBlob = cv2.dnn.blobFromImage(
            img_cv2,
            1.0 / 255,
            (self.inWidth, self.inHeight),
            (0, 0, 0),
            swapRB=False,
            crop=False
        )
        # 将创建的 Blob 设置为神经网络的输入
        self.pose_net.setInput(inpBlob)
        # 指定使用 OpenCV 的 DNN 模块作为后端，并使用 OpenCL 作为目标设备以加速计算。
        self.pose_net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        self.pose_net.setPreferableTarget(cv2.dnn.DNN_TARGET_OPENCL)
        # 前向传播获取  
        # 输出执行前向传播，获取模型的输出结果。output 是一个包含多个热图的张量，每个热图对应一个关键点的概率分布。
        output = self.pose_net.forward()
        # 获取热图的高度 (H) 和宽度 (W)
        H = output.shape[2]
        W = output.shape[3]
        # print(output.shape)

        # 调用 vis_heatmaps 方法可视化热图（这部分代码在当前上下文中未实现具体功能）
        # self.vis_heatmaps(img_cv2, output)

        points = []
        for idx in range(self.num_points):
            probMap = output[0, idx, :, :]  # 置信度图

            # 找到置信度图中的全局最大值
            minVal, prob, minLoc, point = cv2.minMaxLoc(probMap)

            # 将关键点缩放回原始图像大小
            x = (img_width * point[0]) / W
            y = (img_height * point[1]) / H

            if prob > self.threshold:
                points.append((int(x), int(y)))
            else:
                points.append(None)

        return points

    def vis_heatmaps(self, img_cv2, net_outputs):
        """
        可视化热图。

        参数:
            imgfile (str): 输入图像文件路径。
            net_outputs (numpy.ndarray): 网络输出。
        """
        # img_cv2 = cv2.imread(imgfile)
        plt.figure(figsize=[10, 10])

        for pdx in range(self.num_points):
            probMap = net_outputs[0, pdx, :, :]
            probMap = cv2.resize(
                probMap,
                (img_cv2.shape[1], img_cv2.shape[0])
            )
        #     plt.subplot(5, 5, pdx + 1)
        #     plt.imshow(cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB))
        #     plt.imshow(probMap, alpha=0.6)
        #     plt.colorbar()
        #     plt.axis("off")
        # plt.show()

    def vis_pose(self, img_cv2, points):
        """
        可视化姿态估计结果。

        参数:
            imgfile (str): 输入图像文件路径。
            points (list): 关键点坐标列表。
        """
        # img_cv2 = cv2.imread(imgfile)
        # img_cv2_copy = np.copy(img_cv2)

        for idx in range(len(points)):
            if points[idx]:
                cv2.circle(img_cv2,
                           points[idx],
                           8,
                           (0, 255, 255),
                           thickness=-1,
                           lineType=cv2.FILLED)
                cv2.putText(img_cv2,
                            "{}".format(idx),
                            points[idx],
                            cv2.FONT_HERSHEY_SIMPLEX,
                            1,
                            (0, 0, 255),
                            2,
                            lineType=cv2.LINE_AA)

        # 绘制骨架
        for pair in self.point_pairs:
            partA = pair[0]
            partB = pair[1]

            if points[partA] and points[partB]:
                cv2.line(img_cv2,
                         points[partA],
                         points[partB],
                         (0, 255, 255), 3)
                cv2.circle(img_cv2,
                           points[partA],
                           8,
                           (0, 0, 255),
                           thickness=-1,
                           lineType=cv2.FILLED)

        # plt.figure(figsize=[10, 10])
        # plt.subplot(1, 2, 1)
        # plt.imshow(cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB))
        # plt.axis("off")
        # plt.subplot(1, 2, 2)
        # plt.imshow(cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB))
        # plt.axis("off")
        # plt.show()
class general_hand_model(object):
    """
    通用姿态估计模型类，支持多种姿态估计模型（BODY25, COCO, MPI）。
    """

    def __init__(self, modelpath, mode="HAND"):
        """
        初始化姿态估计模型。

        参数:
            modelpath (str): 模型文件路径。
            mode (str): 使用的姿态估计模型类型，默认为 "BODY25"。
        """
        self.inWidth = 368  # 输入图像宽度
        self.inHeight = 368  # 输入图像高度
        self.threshold = 0.1  # 置信度阈值

        if mode == "HAND":
            self.pose_net = self.get_hald_model(modelpath)
    def get_hald_model(self, modelpath):
        """
        加载 MPI 模型。

        参数:
            modelpath (str): 模型文件路径。

        返回:
            mpi_model: 加载的 MPI 模型。
        """
        self.points_name = {
            # "Head": 0, "Neck": 1,
            # "RShoulder": 2, "RElbow": 3, "RWrist": 4,
            # "LShoulder": 5, "LElbow": 6, "LWrist": 7,
            # "RHip": 8, "RKnee": 9, "RAnkle": 10,
            # "LHip": 11, "LKnee": 12, "LAnkle": 13,
            # "Chest": 14, "Background": 15
        }
        self.num_points = 22
        self.point_pairs =  [ [0,1],[1,2],[2,3],[3,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[10,11],[11,12],[0,13],[13,14],[14,15],[15,16],[0,17],[17,18],[18,19],[19,20] ]  # 关键点连接对

        prototxt = os.path.join(
            modelpath,
            "hand/pose_deploy.prototxt"
        )
        caffemodel = os.path.join(
            modelpath,
            "hand/pose_iter_102000.caffemodel"
        )
        hand_model = cv2.dnn.readNetFromCaffe(prototxt, caffemodel)
        return hand_model
    def predict(self,frame):

        t = time.time()
     
        img_height, img_width, _ = frame.shape
        inpBlob = cv2.dnn.blobFromImage(
            frame,
            1.0 / 255,
            (self.inWidth, self.inHeight),
            (0, 0, 0),
            swapRB=False,
            crop=False
        )
        # 设置网络输入
        self.pose_net.setInput(inpBlob)

        # 前向传播获取输出
        output = self.pose_net.forward()
        print("网络处理时间 : {:.3f} 秒".format(time.time() - t))

        # 存储检测到的关键点的空列表
        points = []

        for i in range(self.num_points):
            # 获取对应身体部位的置信度图
            probMap = output[0, i, :, :]
            probMap = cv2.resize(probMap, (img_width, img_height))

            # 找到置信度图中的全局最大值
            minVal, prob, minLoc, point = cv2.minMaxLoc(probMap)

            if prob > self.threshold:
                # 在图像上绘制关键点
                cv2.circle(frame, (int(point[0]), int(point[1])), 8, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
                cv2.putText(frame, "{}".format(i), (int(point[0]), int(point[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, lineType=cv2.LINE_AA)

                # 如果置信度大于阈值，则将点添加到列表中
                points.append((int(point[0]), int(point[1])))
            else:
                points.append(None)
        return points

        

    def vis_hand(self, img_cv2, points):
        """
        可视化姿态估计结果。

        参数:
            imgfile (str): 输入图像文件路径。
            points (list): 关键点坐标列表。
        """
        # img_cv2 = cv2.imread(imgfile)
        # img_cv2_copy = np.copy(img_cv2)

        # 绘制骨架
        for pair in self.point_pairs:
            partA = pair[0]
            partB = pair[1]

            if points[partA] and points[partB]:
                cv2.line(img_cv2, points[partA], points[partB], (0, 255, 255), 2)
                cv2.circle(img_cv2, points[partA], 8, (0, 0, 255), thickness=-1, lineType=cv2.FILLED)
                cv2.circle(img_cv2, points[partB], 8, (0, 0, 255), thickness=-1, lineType=cv2.FILLED)

if __name__ == '__main__':
    print("[INFO] Pose estimation.")

    # img_file = "test.jpg"
    start = time.time()
    modelpath = "./models/"
    # pose_model = general_pose_model(modelpath, mode="COCO")
    hand_model = general_hand_model(modelpath, mode="HAND")
    
    print("[INFO] Model loads time: ", time.time() - start)

    start = time.time()
    rtmp_url="rtmp://127.0.0.1:1935/hls/flow1"
    cap = cv2.VideoCapture(rtmp_url)  # 初始化摄像头
    # if not cap.isOpened():
    #     print("Error: Could not open camera.")
    #     exit()

    while True:
        ret, frame = cap.read()  # 读取摄像头帧
        if not ret:
            print("Error: Could not read frame.")
            break

        start = time.time()
        # res_points = pose_model.predict(frame)
        print("[INFO] Model predicts time: ", time.time() - start)
        res_points1=hand_model.predict(frame)
        # pose_model.vis_pose(frame, res_points)
        hand_model.vis_hand(frame, res_points1)
        print("[INFO] Model vis time: ", frame)
        cv2.imshow('opempose',frame)  # 显示结果
        if cv2.waitKey(1) & 0xFF == ord('q'):  # 按 'q' 键退出
            break

    cap.release()
    cv2.destroyAllWindows()
    # res_points = pose_model.predict(img_file)
    # print("[INFO] Model predicts time: ", time.time() - start)
    # pose_model.vis_pose(img_file, res_points)